import csv
import argparse

def get_execution_units():
    """定义执行单元(EXU)及其包含的功能单元"""
    return {
        "Decode": {
            "description": "Decode Unit",
            "functional_units": {
                "non":    {"number": 1, "pipeline": "pipelined",     "latency": 0},
            }
        },
        "IXU": {
            "description": "Integer Execution Unit",
            "functional_units": {
                "Xu":    {"number": 1, "pipeline": "pipelined",     "latency": 1},
                "Su":    {"number": 1, "pipeline": "pipelined",     "latency": 2},
                "Su-Du": {"number": 1, "pipeline": "non-pipelined", "latency": 19},
            }
        },
        "FXU": {
            "description": "Float Execution Unit",
            "functional_units": {
                "Fadd":   {"number": 1, "pipeline": "pipelined",     "latency": 4},
                "Fmul":   {"number": 1, "pipeline": "pipelined",     "latency": 4},
                "Fdiv-F": {"number": 1, "pipeline": "non-pipelined", "latency": 12},
                "Fdiv-D": {"number": 1, "pipeline": "non-pipelined", "latency": 20},
                "Fcomp":  {"number": 1, "pipeline": "pipelined",     "latency": 2},
                "Fcvt":   {"number": 1, "pipeline": "pipelined",     "latency": 4},
            }
        },
        "VXU-ArithPipe": {
            "description": "Vector Execution Unit, arithmetic pipeline",
            "functional_units": {
                "AddUnit":     {"number": 1, "pipeline": "pipelined", "latency": 1},
                "ShiftUnit":   {"number": 1, "pipeline": "pipelined", "latency": 1},
                "LogicUnit":   {"number": 1, "pipeline": "pipelined", "latency": 1},
                "CompareUnit": {"number": 1, "pipeline": "pipelined", "latency": 1},
                "FmaUnit":     {"number": 1, "pipeline": "pipelined", "latency": 3},
                "FmiscUnit":   {"number": 1, "pipeline": "pipelined", "latency": 2},
                "ImaUnit":     {"number": 1, "pipeline": "pipelined", "latency": 3},
                "CryptoUnit":  {"number": 1, "pipeline": "pipelined", "latency": 1},
                "MatmulUnit":  {"number": 1, "pipeline": "pipelined", "latency": 2},
                "FMatmulUnit": {"number": 1, "pipeline": "pipelined", "latency": 3},
                "SadUnit":     {"number": 1, "pipeline": "pipelined", "latency": 1},
            }
        },
        "VXU-PermPipe": {
            "description": "Vector Execution Unit, permutation pipeline",
            "functional_units": {
                "MoveUnit":     {"number": 1, "pipeline": "pipelined",     "latency": 1},
                "GatherUnit":   {"number": 1, "pipeline": "pipelined",     "latency": 1},
                "Fdiv":         {"number": 1, "pipeline": "non-pipelined", "latency": 32},
                "Idiv":         {"number": 1, "pipeline": "non-pipelined", "latency": 32},
                "SlideUnit":    {"number": 1, "pipeline": "pipelined",     "latency": 1},
                "CompressUnit": {"number": 1, "pipeline": "pipelined",     "latency": 1},
            }
        },
        "VXU-MaskPipe": {
            "description": "Vector Execution Unit, mask pipeline",
            "functional_units": {
                "AddUnit":     {"number": 1, "pipeline": "pipelined", "latency": 1},
                "ShiftUnit":   {"number": 1, "pipeline": "pipelined", "latency": 1},
                "LogicUnit":   {"number": 1, "pipeline": "pipelined", "latency": 1},
                "CompareUnit": {"number": 1, "pipeline": "pipelined", "latency": 1},
                "MaskUnit":    {"number": 1, "pipeline": "pipelined", "latency": 1},
            }
        },
        "LSU": {
            "description": "Load Store Unit",
            "functional_units": {
                "IXU-LSU":              {"number": 3, "pipeline": "pipelined", "latency": 6},
                "VXU-GeneralLoadUnit":  {"number": 1, "pipeline": "pipelined", "latency": 6},
                "VXU-GeneralStoreUnit": {"number": 1, "pipeline": "pipelined", "latency": 6},
                "VXU-IndexLoadUnit":    {"number": 1, "pipeline": "pipelined", "latency": 8},
                "VXU-IndexStoreUnit":   {"number": 1, "pipeline": "pipelined", "latency": 8},
            }
        },
    }

def parse_insn_info(filename):
    """
    解析简单的空格分隔格式
    """
    result = {}

    with open(filename, 'r', encoding='utf-8') as file:
        for line_num, line in enumerate(file, 1):
            line = line.strip()

            if not line:
                continue
            parts = line.split()

            # 检查格式是否正确（应该有4个部分）
            if len(parts) == 4:
                key = parts[0]  # 例如 "rv_zvbb::vandn.vv"
                val1 = parts[1]  # 例如 "ref"
                val2 = parts[2]  # 例如 "ref"
                list_item = parts[3]  # 例如 "zvks"

                # 构建目标格式：("ref", "ref", ["zvks"])
                result[key] = (val1, val2, [list_item])
            else:
                print(f"第 {line_num} 行格式错误，期望4个字段，实际得到 {len(parts)}: {line}")

    return result

def get_instruction_func_mapping():
    """指令映射（包含扩展标签）"""
    data_dict = parse_insn_info('dist/insns_info.txt')
    return data_dict

def search_by_instruction_table():
    """1. 按指令表格方式展示所有指令"""
    exu_info = get_execution_units()
    inst_map = get_instruction_func_mapping()

    print("\n===== 指令表格查询结果 =====")
    print(f"{'指令':<20} | {'扩展标签':<10} | {'EXU':<10} | {'功能单元':<20} | {'流水线类型':<15} | 延迟")
    print("-" * 110)

    seen = set()
    for inst, (exu, func, exts) in sorted(inst_map.items()):
        key = (inst, exu, func)
        if key in seen:
            continue
        seen.add(key)

        if exu not in exu_info:
            print(f"Warning exu is missing: {exu}. info: {inst} {exu} {func} {exts}")
            continue

        func_info = exu_info[exu]["functional_units"][func]
        print(f"{inst:<20} | {', '.join(exts):<10} | {exu:<10} | {func:<20} | {func_info['pipeline']:<15} | {func_info['latency']}")

def search_by_exu():
    """2. 按执行单元(EXU)查询，列出每个EXU包含的指令"""
    exu_info = get_execution_units()
    inst_map = get_instruction_func_mapping()

    # 构建EXU到指令的映射
    exu_to_instructions = {exu: [] for exu in exu_info.keys()}
    seen = set()

    for inst, (exu, func, exts) in inst_map.items():
        key = (inst, exu, func)
        if key in seen:
            continue
        seen.add(key)

        if exu not in exu_info:
            print(f"Warning exu is missing: {exu}. info: {inst} {exu} {func} {exts}")
            continue

        func_info = exu_info[exu]["functional_units"][func]
        exu_to_instructions[exu].append({
            "instruction": inst,
            "func_unit": func,
            "extensions": exts,
            "number": func_info["number"],
            "pipeline": func_info["pipeline"],
            "latency": func_info["latency"]
        })

    print("\n===== 执行单元(EXU)查询结果 =====")
    for exu, details in exu_info.items():
        print(f"\n【EXU: {exu}】({details['description']})")
        print(f"  包含功能单元: {', '.join(details['functional_units'].keys())}")
        print(f"  包含指令:")
        print(f"    {'指令':<20} | {'扩展标签':<10} | {'功能单元':<20} | {'实例数':<6} | {'流水线类型':<15} | 延迟")
        print(f"    {'-' * 80}")

        for item in sorted(exu_to_instructions[exu], key=lambda x: x["instruction"]):
            print(f"    {item['instruction']:<20} | {', '.join(item['extensions']):<10} | {item['func_unit']:<10} | {item['number']:<6} | {item['pipeline']:<15} | {item['latency']}")

def search_by_func_unit():
    """3. 按功能单元(func_unit)查询，列出每个func_unit包含的指令"""
    exu_info = get_execution_units()
    inst_map = get_instruction_func_mapping()

    # 构建功能单元到指令的映射
    func_to_instructions = {}
    # 先初始化所有功能单元
    for exu, details in exu_info.items():
        for func in details["functional_units"].keys():
            func_to_instructions[(exu, func)] = []

    seen = set()
    for inst, (exu, func, exts) in inst_map.items():
        key = (inst, exu, func)
        if key in seen:
            continue
        seen.add(key)

        if exu not in exu_info:
            print(f"Warning exu is missing: {exu}. info: {inst} {exu} {func} {exts}")
            continue

        func_info = exu_info[exu]["functional_units"][func]
        func_to_instructions[(exu, func)].append({
            "instruction": inst,
            "exu": exu,
            "number": func_info["number"],
            "extensions": exts,
            "pipeline": func_info["pipeline"],
            "latency": func_info["latency"]
        })

    print("\n===== 功能单元(func_unit)查询结果 =====")
    for (exu, func), items in sorted(func_to_instructions.items()):
        func_info = exu_info[exu]["functional_units"][func]
        print(f"\n【功能单元: {func}】(所属EXU: {exu})")
        print(f"  实例数量: {func_info['number']} | 流水线类型: {func_info['pipeline']} | 延迟: {func_info['latency']}")
        print(f"  包含指令:")
        print(f"    {'指令':<20} | {'扩展标签':<10} | {'所属EXU':<10}")
        print(f"    {'-' * 50}")

        for item in sorted(items, key=lambda x: x["instruction"]):
            print(f"    {item['instruction']:<20} | {', '.join(item['extensions']):<10} | {item['exu']:<10}")

def export_to_csv(filename="dist/riscv_instruction_details.csv"):
    """导出所有指令数据到CSV"""
    exu_info = get_execution_units()
    inst_map = get_instruction_func_mapping()

    with open(filename, mode='w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow(["指令", "扩展标签", "执行单元(EXU)", "功能单元", "流水线类型", "延迟(周期)"])

        seen = set()
        for inst, (exu, func, exts) in sorted(inst_map.items()):
            key = (inst, exu, func)
            if key in seen:
                continue
            seen.add(key)

            if exu not in exu_info:
                print(f"Warning exu is missing: {exu}. info: {inst} {exu} {func} {exts}")
                continue

            func_info = exu_info[exu]["functional_units"][func]
            writer.writerow([
                inst,
                ", ".join(exts),
                exu,
                func,
                func_info["pipeline"],
                func_info["latency"]
            ])
    print(f"\n数据已导出到 {filename}")

def export_execution_units_to_csv(filename="dist/riscv_execution_units.csv"):
    """将execution_units表格导出为CSV文件"""
    exu_info = get_execution_units()

    with open(filename, mode='w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        # 表头：执行单元、执行单元描述、功能单元、功能单元实例数、流水线类型、延迟
        writer.writerow(["执行单元(EXU)", "功能单元(func_unit)", "实例数量(number)", "流水线类型(pipeline)", "延迟(latency)", "EXU描述"])

        # 遍历所有执行单元及其功能单元
        for exu, exu_details in exu_info.items():
            exu_desc = exu_details["description"]
            for func_unit, func_details in exu_details["functional_units"].items():
                writer.writerow([
                    exu,
                    func_unit,
                    func_details["number"],
                    func_details["pipeline"],
                    func_details["latency"],
                    exu_desc,
                ])

    print(f"执行单元(execution_units)数据已导出到 {filename}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--interaction', help='交互模式', action='store_true')

    args = parser.parse_args()

    if args.interaction:
        while True:
            print("\n===== RISC-V指令查询工具 =====")
            print("1. 按指令表格方式展示所有指令")
            print("2. 按执行单元(EXU)查询指令")
            print("3. 按功能单元(func_unit)查询指令")
            print("4. 导出数据到CSV文件")
            print("5. 导出执行单元(execution_units)到CSV文件")
            print("0. 退出")

            choice = input("请选择操作(0-5): ")
            if choice == "1":
                search_by_instruction_table()
            elif choice == "2":
                search_by_exu()
            elif choice == "3":
                search_by_func_unit()
            elif choice == "4":
                export_to_csv()
            elif choice == "5":
                export_execution_units_to_csv()
            elif choice == "0":
                print("退出程序")
                break
            elif choice == "q":
                print("退出程序")
                break
            else:
                print("无效选择，请重试")
    else:
        export_to_csv()
        export_execution_units_to_csv()
        print("\n===== RISC-V指令查询工具 =====")
        print("1. 导出数据到CSV文件")
        print("2. 导出执行单元(execution_units)到CSV文件")

if __name__ == "__main__":
    main()
